{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "daed6e8f-8426-4d44-9d48-5cda928f82b8",
   "metadata": {},
   "source": [
    "# Complete Document Retrieval and RAG\n",
    "\n",
    "## Overview\n",
    "- **Retrieval Pipeline** With customers having the ability to enter any number of possibilities into the solution, it is helpful to detect intent and normalize the query. Few-shots are a useful tool to tailor the normalization to the nature of the query in-line. \n",
    "- **Advanced methods** For more complex cases, it can be beneficial to generate hypothetical queries and documents solving for sub-queries and improving the semantic similarity.\n",
    "- **Model answer generation** Once the model is shown a set of documents, it must generate an answer while staying as closely aligned to the contents of the documents as possible. We cover self-verification and citation as methods giving greater flexibility to the model for a given query and set of retrieved documents.\n",
    "\n",
    "## Context\n",
    "\n",
    "Retrieval Augmented Generation (RAG) requires the indexation of relevant unstructured documents into a vector database. Then given a customer query, the relevant are retrieved and past as context to the model, which generates an answer. This can best be described by the following flow.\n",
    "\n",
    "<img src=\"./images/rag-architecture.png\" />\n",
    "\n",
    "Once our documents (PDFs, CSV, Tables, JSON, ...) have been indexed into our knowledge base, we start working towards retrieval of a relevant subset of documents based on a given query. For many applications, the success of the retrieval is a strong indicator for the performance of the overall response. This notebook assumes you are familiar with the basics of RAG, embedding models and vector databases.\n",
    "\n",
    "In this notebook, we seek to go beyond RAG to generate the model answer by applying other relevant steps in the answer pipeline.\n",
    "\n",
    "<h2>Prerequisites</h2>\n",
    "\n",
    "Before you can use Amazon Bedrock, you must carry out the following steps:\n",
    "\n",
    "- Sign up for an AWS account (if you don't already have one) and IAM Role with the necessary permissions for Amazon Bedrock, see [AWS Account and IAM Role](https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html#new-to-aws){:target=\"_blank\"}.\n",
    "- Request access to the foundation models (FM) that you want to use, see [Request access to FMs](https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html#getting-started-model-access){:target=\"_blank\"}. \n",
    "    \n",
    "<h2>Setup</h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f9239b9-3894-449c-97e8-7b04c680817c",
   "metadata": {},
   "source": [
    "We import the relevant objects used in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7a8e04e-ddba-491d-a52a-448138bf07f8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:03:08.288566Z",
     "start_time": "2025-02-13T20:03:08.281477Z"
    }
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import faiss\n",
    "import re\n",
    "from operator import itemgetter\n",
    "from typing import List\n",
    "from langchain_aws.chat_models.bedrock import ChatBedrock\n",
    "from langchain_aws import BedrockEmbeddings\n",
    "from langchain_core.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    FewShotChatMessagePromptTemplate,\n",
    ")\n",
    "from IPython.display import display_markdown, Markdown\n",
    "from langchain_community.docstore import InMemoryDocstore\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from pydantic import BaseModel, Field\n",
    "from langchain_core.example_selectors import SemanticSimilarityExampleSelector\n",
    "from langchain_core.prompts import HumanMessagePromptTemplate, AIMessagePromptTemplate\n",
    "from langchain.output_parsers import PydanticToolsParser\n",
    "from langchain_community.retrievers import WikipediaRetriever\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.runnables import (\n",
    "    RunnableLambda,\n",
    "    RunnableParallel,\n",
    "    RunnablePassthrough,\n",
    "    RunnableBranch,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8141160-6799-4c30-8cf5-3f917c87a646",
   "metadata": {},
   "source": [
    "Although this example leverages Nova Pro & Nova Lite, Bedrock supports many other models. This full list of models and supported features can be found [here](https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html). The models are invoked via `bedrock-runtime`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3c95bebf-c30b-47a3-8566-438275fe37da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:02:10.404760Z",
     "start_time": "2025-02-13T20:02:10.284615Z"
    }
   },
   "outputs": [],
   "source": [
    "region = 'us-west-2'\n",
    "bedrock = boto3.client(\n",
    "    service_name = 'bedrock-runtime',\n",
    "    region_name = region,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "144041e6-f221-4f3a-a546-5356a4edee29",
   "metadata": {},
   "source": [
    "We use `ChatBedrock` and `BedrockEmbeddings` to interact with the Bedrock API. We enable `beta_use_converse_api` to use the Converse API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1723192-9fed-4fb1-8dc0-9a9d26f531f8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:12.275860Z",
     "start_time": "2025-02-13T20:20:12.273354Z"
    }
   },
   "outputs": [],
   "source": [
    "modelId = \"us.amazon.nova-lite-v1:0\"\n",
    "nova = ChatBedrock(\n",
    "    model_id=modelId,\n",
    "    client=bedrock,\n",
    "    beta_use_converse_api=True\n",
    ")\n",
    "embeddingId = \"amazon.titan-embed-text-v2:0\"\n",
    "embeddings = BedrockEmbeddings(\n",
    "    model_id=embeddingId,\n",
    "    client=bedrock)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc2bb343-bf2c-4429-bc20-90e4dc48bcd4",
   "metadata": {},
   "source": [
    "We correctly get a generic answer message from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e7e9f73a-36bc-4767-a94a-3ca85c90e8ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:17.121473Z",
     "start_time": "2025-02-13T20:20:13.177172Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Of course! I'd be happy to help you with your travel needs today. Here are some key areas we can cover:\n",
       "\n",
       "### 1. **Packing:**\n",
       "   - **Clothing:** Pack versatile clothing that can be mixed and matched. Don’t forget to include a few items that match the local climate and any specific activities you plan to do.\n",
       "   - **Essentials:** Don’t forget essentials like your passport, tickets, money, and any necessary documents.\n",
       "   - **Tech:** Make sure to pack your phone, charger, adapters, headphones, and any other electronics you need.\n",
       "   - **Toiletries:** Pack a small toiletry bag with essentials like toothbrush, toothpaste, deodorant, sunscreen, and any medications you might need.\n",
       "\n",
       "### 2. **Transportation:**\n",
       "   - **Airport Transfers:** If you need to get to the airport, consider booking a taxi, rideshare, or shuttle service in advance.\n",
       "   - **Public Transport:** Familiarize yourself with the local public transport options such as buses, trains, or subways.\n",
       "   - **Car Rentals:** If you’re renting a car, make sure you have your driver’s license, rental agreement, and any necessary insurance.\n",
       "\n",
       "### 3. **Accommodation:**\n",
       "   - **Booking:** Ensure your hotel or accommodation booking is confirmed. Have the reservation details handy.\n",
       "   - **Check-In:** Know the check-in time and what documents you might need to present.\n",
       "\n",
       "### 4. **Local Transportation:**\n",
       "   - **Navigation:** Download maps and navigation apps. Consider apps like Google Maps, Citymapper, or local transit apps.\n",
       "   - **Public Transport Tickets:** Purchase any necessary public transport tickets or passes before you need them.\n",
       "\n",
       "### 5. **Dining:**\n",
       "   - **Local Cuisine:** Research popular local dishes and restaurants. Consider making reservations if necessary.\n",
       "   - **Food Safety:** Be mindful of local food safety practices to avoid any digestive issues.\n",
       "\n",
       "### 6. **Communication:**\n",
       "   - **Local SIM Card:** If you need mobile data, consider getting a local SIM card or an international plan.\n",
       "   - **Translation Apps:** Download translation apps to help with language barriers.\n",
       "\n",
       "### 7. **Health and Safety:**\n",
       "   - **Travel Insurance:** Ensure you have travel insurance that covers medical emergencies, trip cancellations, and lost belongings.\n",
       "   - **Health Precautions:** Check if there are any health advisories or required vaccinations for your destination.\n",
       "\n",
       "### 8. **Activities and Sightseeing:**\n",
       "   - **Itinerary:** Plan your itinerary but remain flexible. Note opening hours and any special events.\n",
       "   - **Tickets:** Purchase tickets in advance for popular attractions to avoid long lines.\n",
       "\n",
       "### 9. **Cultural Etiquette:**\n",
       "   - **Local Customs:** Research local customs and etiquette to ensure respectful and enjoyable interactions.\n",
       "\n",
       "### 10. **Emergency Contacts:**\n",
       "   - **Local Emergency Numbers:** Know the local emergency numbers and the address of your country’s embassy or consulate.\n",
       "   - **Contact List:** Keep a list of important contacts, including local friends, family, and hotel staff.\n",
       "\n",
       "If you have specific questions or need more detailed information on any of these topics, feel free to ask! Safe travels!"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(nova.invoke(\"Help me with my travel needs today.\").content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b8c374c-2c9e-4093-ab7c-48b9c39ff3c3",
   "metadata": {},
   "source": [
    "## Reformating the initial query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac197d7a-0397-4ab8-8f82-3d8dc63fdc3a",
   "metadata": {},
   "source": [
    "### Intent Detection\n",
    "\n",
    "In order to limit the scope of answers handled by the solution with RAG, a common first step in the answer pipeline is **Intent Detection or Classification**. This step is important to ensure the relevancy of the question to the indexed content, which works to limit the model's tendancy to answer questions that may not have been accounted for or tested by the application developers.\n",
    "\n",
    "When requesting some information that is irrelevant to the previously stated purpose, we quickly see the model attempting to provide an answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c4f19692-d210-45c0-81fb-d50c3f9ae1fb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:32.732386Z",
     "start_time": "2025-02-13T20:20:22.853558Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Sure, I'd be happy to help you learn more about your mom's pie recipe! Here are some questions that might help you gather more information:\n",
       "\n",
       "1. What type of pie is it? Sweet or savory? Fruit, custard, cream, or something else?\n",
       "2. What are the main ingredients? Flour, sugar, butter, eggs, and any other key components?\n",
       "3. Are there any unique or unusual ingredients or flavors in the recipe?\n",
       "4. How is the crust prepared? Is it a traditional pie crust, or is there something different about it?\n",
       "5. What is the baking process like? What temperature and for how long is the pie baked?\n",
       "6. Are there any special tips or tricks your mom uses to make the pie turn out perfectly?\n",
       "\n",
       "Once you have some more information about the recipe, we can start to explore ways to improve it, substitute ingredients, or make variations. If you have specific questions about any part of the recipe, feel free to ask!"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(nova.invoke(\"I want to learn more about my mom's pie recipe\").content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c00d6258-510c-46ee-b5ae-2f31800be497",
   "metadata": {},
   "source": [
    "Hence, we provide an initial system prompt defining the model's role as an intent classifier. We supply the classes and few-shots to improve performance and ensure the model is aligned to the desired intended output, which needs to include `<intention></intention>` tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e16103d7-2c42-449a-b50b-bb20631dceea",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:34.495333Z",
     "start_time": "2025-02-13T20:20:34.493185Z"
    }
   },
   "outputs": [],
   "source": [
    "intent_system_prompt = \"\"\"You are a precise classifier. Your task is to assess customer intent and categorize customer inquiry into one of the intentions. \n",
    "\n",
    "Intentions with their description:\n",
    "vacation: Information on vacations, various travel destinations and my recent travels.\n",
    "contact: Expressing the desire to talk to support.\n",
    "irrelevant: Not related to vacations and travel.\n",
    "\n",
    "Here is an example of how to respond in a standard interaction:\n",
    "<example>\n",
    "    Human: I am seeking a place that is sunny a family friendly.\n",
    "    AI: <intention>vacation</intention>\n",
    "</example>\n",
    "<example>\n",
    "    Human: I want to learn more about my mom's pie recipe\n",
    "    AI: <intention>irrelevant</intention>\n",
    "</example>\n",
    "<example>\n",
    "    Human: I want to talk to a someone.\n",
    "    AI: <intention>contact</intention>\n",
    "</example>\n",
    "\n",
    "Think about your answer first before you respond. Think step-by-step and insert the classification in <intention></intention> tags and do not include anything after.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f456a5-cc22-4235-9433-406487284269",
   "metadata": {},
   "source": [
    "We supply the prompt as part of `ChatPromptTemplate`and use the pipe operator to define a chain connecting the model to the resulting prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c124c190-f1b3-4cfa-89c7-c7c409bd0411",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:36.598760Z",
     "start_time": "2025-02-13T20:20:36.596229Z"
    }
   },
   "outputs": [],
   "source": [
    "intent_detection_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", intent_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "intent_detection_chain = intent_detection_prompt | nova"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bde5c4a1-c682-4632-9b79-0503a4bd5f10",
   "metadata": {},
   "source": [
    "We invoke the model with the same query and notice the classification result. We invite you to try additional questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8fce706b-187f-4518-9927-c72ae44f3ad9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:20:45.684811Z",
     "start_time": "2025-02-13T20:20:45.312662Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<intention>irrelevant</intention>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(intent_detection_chain.invoke(\"Tell me about my mother's pie recipe\").content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0efc0d-9826-4a54-94e6-81056a58c166",
   "metadata": {},
   "source": [
    "Since we expect the answer to always contain these tags, we can parse it and branch off depending on the model's classification. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8804bd6e-4c66-458e-98c2-7f4948bad4fa",
   "metadata": {},
   "source": [
    "### Dynamic few-shots\n",
    "\n",
    "Although static few-shots are helpful, they have two major obstacles. On the one hand, they do not cover the breadth of necessary examples, and on the other, given that any submitted query is rarely relevant to all supplied examples, they often introduce unecessary tokens and noise to the prompt. In constrast, supplying dynamic few-shots from a larger corpus of examples enables us to select a number of the most relevant examples prior to inference. Evidently, these are determined by the nature of the query. Although we apply it to intend classification, dynamic few-shots can be applied anywhere in the RAG pipeline and generally yield stronger results compared to static examples. \n",
    "\n",
    "We bootstrap `few_shot_library` using examples distilled by **Amazon Nova Pro**. It is important to continuously iterate on the library after the initial deployment. During this phase, it is a general best practice to collect and label real interactions where the model made mistakes and append those to the set of examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "61480afe-7501-4e3b-a191-7d66dd456964",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:08.620638Z",
     "start_time": "2025-02-13T20:21:08.615379Z"
    }
   },
   "outputs": [],
   "source": [
    "few_shot_library = [\n",
    "    {\n",
    "        \"question\": \"Can you recommend some tropical beach destinations?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I need to speak with a customer service representative.\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the best way to cook spaghetti?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Are there any family-friendly resorts in Florida?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do I file a complaint about my recent stay?\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the weather like in Paris in June?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Can you help me with my car insurance claim?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I'd like to book an all-inclusive Caribbean cruise.\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Is there a phone number for your reservations team?\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the best way to learn a new language?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Are there any good hiking trails in Yellowstone?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I need to update my billing information.\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do I make homemade bread?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What are some popular tourist attractions in Rome?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Can I speak with a manager about my recent experience?\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the best time to visit Japan?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do I reset my Netflix password?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Are there any good ski resorts in Colorado?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I need help with my online booking.\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the plot of the latest Marvel movie?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Can you suggest some budget-friendly European cities?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do I request a refund for my canceled trip?\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the best way to train a puppy?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Are there any good wildlife safaris in Africa?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I need to change my flight reservation.\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What are some must-see landmarks in New York City?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do I fix a leaky faucet?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Can you recommend some romantic getaways for couples?\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"I have a question about my loyalty points balance.\",\n",
    "        \"class\": \"contact\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What's the best way to prepare for a job interview?\",\n",
    "        \"class\": \"irrelevant\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Tell me about my travel history\",\n",
    "        \"class\": \"vacation\"\n",
    "    },\n",
    "    \n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fec08d8-4273-4689-9a5c-b189186a451d",
   "metadata": {},
   "source": [
    "In this notebook, we use FAISS (Facebook AI Similarity Search) [(github)](https://github.com/facebookresearch/faiss), which is an open-source library developed by Facebook AI Research for efficient similarity search and clustering of dense vector embeddings. We call the Lanchain's `FAISS` object to interact with the in-memory vector store.\n",
    "\n",
    "We embed the examples using the Titan Embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "87fce4d9-b019-4f22-a0fc-bf15c680932c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:10.159374Z",
     "start_time": "2025-02-13T20:21:10.157315Z"
    }
   },
   "outputs": [],
   "source": [
    "embedding_size = 1536\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "\n",
    "vectorstore = FAISS(embeddings, index, InMemoryDocstore({}), {})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec04c8c5-21dc-4c89-bb7d-dd0af1e91b0c",
   "metadata": {},
   "source": [
    "We use `SemanticSimilarityExampleSelector` to dynamically select the `k` most relevant examples based on our query. When instantiated, this object embeds the set of examples into our vector store of choice. `FewShotChatMessagePromptTemplate` defines the formatting of the selected examples into a given prompt. We define the template to be consistent with what will be generated by the model during intent classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "303a35b5-a45e-42ce-ad2e-e0043f204c6c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:20.266966Z",
     "start_time": "2025-02-13T20:21:15.168799Z"
    }
   },
   "outputs": [],
   "source": [
    "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    few_shot_library,\n",
    "    embeddings,\n",
    "    vectorstore,\n",
    "    k=5,\n",
    ")\n",
    "\n",
    "few_shot_prompt = FewShotChatMessagePromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=(\n",
    "        HumanMessagePromptTemplate.from_template(\"{question}\")\n",
    "        + AIMessagePromptTemplate.from_template(\"<intention>{class}</intention>\")\n",
    "    ),\n",
    "    input_variables=[\"question\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5038eca0-b90b-4847-997b-4fd84b17c082",
   "metadata": {},
   "source": [
    "We print the relevant examples for a given query. Notice that the distribution of labels will change based on the nature of the query. This helps further align the model with our expectations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8c171e51-e5cd-48b6-887b-09c915672607",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:20.450437Z",
     "start_time": "2025-02-13T20:21:20.283536Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Human: Tell me about my travel history\n",
       "AI: <intention>vacation</intention>\n",
       "Human: I'd like to book an all-inclusive Caribbean cruise.\n",
       "AI: <intention>vacation</intention>\n",
       "Human: Can you suggest some budget-friendly European cities?\n",
       "AI: <intention>vacation</intention>\n",
       "Human: Can I speak with a manager about my recent experience?\n",
       "AI: <intention>contact</intention>\n",
       "Human: How do I request a refund for my canceled trip?\n",
       "AI: <intention>contact</intention>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(few_shot_prompt.format(question=\"tell me about my travels\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e6abc8-0bea-42fb-8925-e33bbaa6a6f3",
   "metadata": {},
   "source": [
    "We redefine the system prompt to accomodate for the dynamic few-shots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "db144844-194f-4c8b-b623-23c0d901c276",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:20.478221Z",
     "start_time": "2025-02-13T20:21:20.476440Z"
    }
   },
   "outputs": [],
   "source": [
    "few_shot_intent_system_prompt = \"\"\"You are a precise classifier. Your task is to assess customer intent and categorize customer inquiry into one of the intentions. \n",
    "\n",
    "Intentions with their description:\n",
    "vacation: Information on vacations, various travel destinations and my recent travels.\n",
    "contact: Expressing the desire to talk to support.\n",
    "irrelevant: Not related to vacations and travel.\n",
    "\n",
    "Here is an example of how to respond in a standard interaction:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b6c2811-3a7e-4d8b-af4b-e8b326ed6b74",
   "metadata": {},
   "source": [
    "We redefine the prompt template to accomodate for the dynamic few-shots. As expected, the final string created from `intent_detection_prompt` will change based on message similarity to previous examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e215997a-e940-4a94-a4f3-a424b595c90c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:20.501123Z",
     "start_time": "2025-02-13T20:21:20.498866Z"
    }
   },
   "outputs": [],
   "source": [
    "few_shot_intent_detection_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", few_shot_intent_system_prompt),\n",
    "        few_shot_prompt,\n",
    "        (\"human\", \"Think step-by-step and always ensure you insert the classification in <intention></intention> tags and do not include anything after.\\\n",
    "        Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "few_shot_intent_chain = intent_detection_prompt | nova"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efde4eed-db7c-40aa-a037-fbbd30ba0492",
   "metadata": {},
   "source": [
    "We test the newly created chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "94b35d0b-1fbe-4a0b-a331-b445d0388607",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:22.992234Z",
     "start_time": "2025-02-13T20:21:22.564561Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<intention>vacation</intention>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(few_shot_intent_chain.invoke({\"question\": \"tell me about my travel history\"}).content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a359d7ba-ce4d-421f-901b-1c97aab0590c",
   "metadata": {},
   "source": [
    "### Normalizing the user message\n",
    "\n",
    "We may want to restrict the queries that are sent to downstream inference without restricting the user experience. Normalizing messages enables us to do exactly this. It can often be used to set a certain tone, reduce length and extract the specific purpose of the message while reducing unecessary noise. Notice the role the rule book plays in determining the nature of the returned message.\n",
    "\n",
    "Alternatively, it is common to supply few-shot examples as we have done in the previous step. We again return the resulting message in between tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3018bf44-5cb6-48bd-9846-b9bf08562e85",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:28.280127Z",
     "start_time": "2025-02-13T20:21:28.278111Z"
    }
   },
   "outputs": [],
   "source": [
    "norm_system_prompt = \"\"\"You are a precise message synthesizer. Your task is to write a condensed message encompassing the latest original message's intent and main keywords. \n",
    "The condensed message must follow the rule book.\n",
    "\n",
    "Rule book:\n",
    "- Must be a complete sentence formulated as a request from the perspective of the original requester.\n",
    "- No longer than 2 short sentences with no concatination.\n",
    "- Never include names.\n",
    "- It is safe to reformulate questions with only keyword as looking for information on the place they mention.\n",
    " \n",
    "Think about your answer first before you respond. Think step-by-step and the condensed message in <condensed_message></condensed message> tags and do not include anything after.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ef0927d-2ca2-40c7-8068-180f7f5eac6f",
   "metadata": {},
   "source": [
    "We define the prompt template incorporating the system prompt with the user defined message. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5fe2af57-4dbe-4e1b-bf68-7a59d72eb5d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:21:29.876339Z",
     "start_time": "2025-02-13T20:21:29.874433Z"
    }
   },
   "outputs": [],
   "source": [
    "norm_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", norm_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "norm_chain = norm_prompt | nova"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbad2627-6be4-41af-8520-de8879165732",
   "metadata": {},
   "source": [
    "When executing the chain on a longer query, the returned message pulls out only the information necessary to the task at hand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f9467956-8af5-4acc-9197-9a0c34428263",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:07:09.334143Z",
     "start_time": "2025-02-13T20:07:08.647094Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<condensed_message>Could you provide information on my travel history?</condensed_message>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(norm_chain.invoke({\"question\": \"\"\"I have been all around the world seing a bunch of stuff. \n",
    "I met a bunch of people like Bernard and Tamy. Tell me about my travel history\"\"\"}).content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80d22a61-b1f2-4de1-9702-214b0c70ebf0",
   "metadata": {},
   "source": [
    "When executing the chain on a query that only has keywords, the model fills in the gap to provide additional context. Although the initial queries are quite different, notice that their resulting output is quite similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "95585830-965a-4bbf-9fef-7304a74b48c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:07:45.094986Z",
     "start_time": "2025-02-13T20:07:44.531820Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<condensed_message>Could you provide information on New York?</condensed_message>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_markdown(Markdown(norm_chain.invoke({\"question\": \"\"\"New York\"\"\"}).content))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acf9ab56-f876-43f0-a9c7-a2d84449e8fb",
   "metadata": {},
   "source": [
    "Once we have detected the message's intent and normalized it to some extent, we are able to have much greater assurance as to the nature of the messages sent to subsequent steps, namely the retrieval."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "565f4fb5-c973-42cb-b8ca-64516fcb12c3",
   "metadata": {},
   "source": [
    "## Advanced methods of retrieval\n",
    "\n",
    "The main driver of performance for RAG pipelines is the retrieval mechanism. This step involves identifying a subset of documents that are most relevant to the original query. The common baseline is generally to embed the query in its original form and pull the top-K nearest documents. However, for some datasets this begins to fall short in cases where queries address multiple topics or, more generally, are phrased in a way that is incompatible or is dissimilar to the documents that should be retrieved. We look at how it is possible to improve on these types of queries. \n",
    "\n",
    "Given the increase complexity of the tasks in this section, we choose to leverage Amazon Nova Pro in this part of the pipeline. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "85f16644-77cf-4e4a-bc0f-bf5e6d1eade8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:09:41.919318Z",
     "start_time": "2025-02-13T20:09:41.912616Z"
    }
   },
   "outputs": [],
   "source": [
    "modelId = \"us.amazon.nova-pro-v1:0\"\n",
    "nova_pro = ChatBedrock(\n",
    "    model_id=modelId,\n",
    "    client=bedrock,\n",
    "    beta_use_converse_api=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f9cdc5-5b83-4be3-be7e-32c96aebc71d",
   "metadata": {},
   "source": [
    "### Decomposition\n",
    "\n",
    "For more complex queries, it may be helpful to breakdown the original question into sub-problems each having their own retrieval step. We perform query decomposition to return the original question or an equivalent set of questions each with a single target.\n",
    "\n",
    "This process is driven by the underlying model. We define the system prompt describing the intended task and supply static few-shot examples to enable the model to better generalize. Removing these examples yields results that are less robust."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f5ab2ee7-9785-45ab-a149-6c45b3b2aa69",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:09:42.674598Z",
     "start_time": "2025-02-13T20:09:42.672977Z"
    }
   },
   "outputs": [],
   "source": [
    "decomp_system_prompt = \"\"\"You are a expert assistant that prepares queries that will be sent to a search component. \n",
    "These queries may be very complex. Your job is to simplify complex queries into multiple queries that can be answered in isolation to eachother.\n",
    "\n",
    "If the query is simple, then keep it as it is.\n",
    "\n",
    "If there are acronyms or words you are not familiar with, do not try to rephrase them.\n",
    "Here is an example of how to respond in a standard interaction:\n",
    "<example>\n",
    "- Query: Did Meta or Nvidia make more money last year?\n",
    "Decomposed Questions: [SubQuery(sub_query='How much profit did Meta make last year?'), SubQuery(sub_query'How much profit did Nvidia make last year?')]\n",
    "</example>\n",
    "<example>\n",
    "- Query: What is the capital of France?\n",
    "Decomposed Questions: [SubQuery(sub_query='What is the capital of France?')]\n",
    "</example>\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7abe7d9f-d645-4a8f-9139-ff87a5a9cd2e",
   "metadata": {},
   "source": [
    "To ensure a consistent format is returned for subsequent steps, we use Pydantic, a data-validation library. We rely on a Pydantic-based helper function for doing the tool config translation for us in a way that ensures we avoid potential mistakes when defining our tool config schema in a JSON dictionary.\n",
    "\n",
    "We define `SubQuery` to be a query corresponding to a subset of the points of a larger parent query. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "13e37172-db97-44d4-91d7-0f70504be87f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:09:43.492896Z",
     "start_time": "2025-02-13T20:09:43.490253Z"
    }
   },
   "outputs": [],
   "source": [
    "class SubQuery(BaseModel):\n",
    "    \"\"\"You have performed query decomposition to generate a subquery of a question\"\"\"\n",
    "\n",
    "    sub_query: str = Field(description=\"A unique subquery of the original question.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1e910b8-2271-4ea8-bc03-a1bdbd37b8ec",
   "metadata": {},
   "source": [
    "We define the prompt template leveraging the previously defined system prompt. We then expose `SubQuery` as a tool the model can leverage. This enables to model to format one or more requests to this tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1cf57154-3307-407a-844f-41576bfd0b64",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:09:45.239217Z",
     "start_time": "2025-02-13T20:09:45.036928Z"
    }
   },
   "outputs": [],
   "source": [
    "query_decomposition_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", decomp_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "llm_with_tools = nova_pro.bind_tools([SubQuery])\n",
    "decomp_query_analyzer = query_decomposition_prompt | llm_with_tools | PydanticToolsParser(tools=[SubQuery])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb9173ba-2af7-4dc8-a708-8d32db799fc5",
   "metadata": {},
   "source": [
    "We asking a broad question about multiple destinations, the model chooses to return multiple calls to `SubQuery`. Each can be sent for document retrieval in parallel, thus ensuring we do not encure additional latency beyond that of the model inferencing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f7566eca-fc92-41ff-ba5f-ae588540217e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:10:42.431837Z",
     "start_time": "2025-02-13T20:10:40.935080Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SubQuery(sub_query='How do you plan a vacation in Thailand?'),\n",
       " SubQuery(sub_query='How do you plan a vacation in California?')]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "queries = decomp_query_analyzer.invoke({\"question\": \"How do go on vacation in thailand and in California?\"})\n",
    "queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec0ab1a-90ea-473c-92e7-4c5fb4108b57",
   "metadata": {},
   "source": [
    "### Expansion\n",
    "\n",
    "Query expansion is similar to decomposition in that it produces multiple queries as a strategy to improve the odds of hitting a relevant result. However, expansion returns multiple different wordings of the original query.  \n",
    "\n",
    "We define the system prompt to consistently return 3 versions of the original query. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4cbcd9b3-a6fb-456f-88ba-be9eb40d240f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:10:44.839343Z",
     "start_time": "2025-02-13T20:10:44.837204Z"
    }
   },
   "outputs": [],
   "source": [
    "paraphrase_system_prompt = \"\"\"You are an expert at converting user questions into database queries. \n",
    "You have access to a database of travel destinations and a list of recent destinations for travelers. \n",
    "\n",
    "Perform query expansion. If there are multiple common ways of phrasing a user question \n",
    "or common synonyms for key words in the question, make sure to return multiple versions \n",
    "of the query with the different phrasings.\n",
    "\n",
    "If there are acronyms or words you are not familiar with, do not try to rephrase them.\n",
    "\n",
    "Always return at least 3 versions of the question.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc546c0b-64eb-4463-a93b-4b5a3c6a3bd7",
   "metadata": {},
   "source": [
    "We define the prompt template leveraging the previously defined system prompt. We then expose `ParaphrasedQuery` as a tool the model can leverage. This enables to model to format one or more requests to this tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5676aca2-d322-4dbe-bc7d-a3ab5ece2796",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:10:46.572300Z",
     "start_time": "2025-02-13T20:10:46.568609Z"
    }
   },
   "outputs": [],
   "source": [
    "class ParaphrasedQuery(BaseModel):\n",
    "    \"\"\"You have performed query expansion to generate a paraphrasing of a question.\"\"\"\n",
    "\n",
    "    paraphrased_query: str = Field(description=\"A unique paraphrasing of the original question.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "219d3baf-47ae-45db-838e-46bcb89071fd",
   "metadata": {},
   "source": [
    "We define the prompt template leveraging the previously defined system prompt. We then expose `ParaphrasedQuery` as a tool the model can leverage. This enables to model to format one or more requests to this tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "61071899-b1d1-42e4-a91f-dbf3e5bb5859",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:10:54.156134Z",
     "start_time": "2025-02-13T20:10:54.089587Z"
    }
   },
   "outputs": [],
   "source": [
    "query_expansion_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", paraphrase_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "llm_with_tools = nova_pro.bind_tools([ParaphrasedQuery])\n",
    "query_expansion = query_expansion_prompt | llm_with_tools | PydanticToolsParser(tools=[ParaphrasedQuery])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39c037c-6b4f-4768-943b-19e9dd023b16",
   "metadata": {},
   "source": [
    "Now no matter the nature of the query, the model generates alternatives that can be sent for retrieval in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c5470a22-899a-4bb9-8224-ce9c3cd2a859",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:10:57.646048Z",
     "start_time": "2025-02-13T20:10:55.661343Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[ParaphrasedQuery(paraphrased_query='what are the steps to travel to Canada and Mexico?'),\n",
       " ParaphrasedQuery(paraphrased_query='what is the process for traveling to Canada and Mexico?'),\n",
       " ParaphrasedQuery(paraphrased_query='how can I travel to Canada and Mexico?')]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_expansion.invoke({\"question\": \"how to use travel to Canada and to Mexico?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c56aaaf5-ac64-4e58-92a9-898be2174e91",
   "metadata": {},
   "source": [
    "### Hypothetical Document Embeddings (HyDE)\n",
    "\n",
    "Given that models have been trained large volumes of data, we can generate a relevant hypothetical document to answer the user question. Then for retrieval, this new (or *hypethetical*) document can be embedded with the original query. This approach has been shown in [Precise Zero-Shot Dense Retrieval without Relevance Labels](https://arxiv.org/abs/2212.10496) to improve recall. We define the system prompt relevant to this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5ed9886a-707c-4782-bdc1-0239f54134b1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:11:00.073889Z",
     "start_time": "2025-02-13T20:11:00.072316Z"
    }
   },
   "outputs": [],
   "source": [
    "hyde_system_prompt = \"\"\"You are an expert about travel destinations all over the worlds. Your task is to provide your best response based on the question.\n",
    "You need to produce a high-quality and complete sentence hyper focused on answer the question. \n",
    "Do not answer in bulletpoints.\n",
    "\n",
    "Think about your answer first before you respond. Think step-by-step and the answer in <hyde></hyde> tags and do not include anything after.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f2de46d-651b-4b7f-af83-273966a03478",
   "metadata": {},
   "source": [
    "We define the prompt template leveraging the previously defined system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f741b5a3-481c-467b-b454-40d243232747",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:11:05.606832Z",
     "start_time": "2025-02-13T20:11:05.603928Z"
    }
   },
   "outputs": [],
   "source": [
    "hyde_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", hyde_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "hyde_chain = hyde_prompt | nova_pro | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f319385-8b34-489b-84d9-8ec921e4853c",
   "metadata": {},
   "source": [
    "We produce a document for the query in between tags that is be appended at retrieval time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4a2467e3-94d9-461d-9a76-c3f8bb7862a2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:11:21.309801Z",
     "start_time": "2025-02-13T20:11:17.727658Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<hyde>To answer the customer's question about going on vacation in Thailand and California, I need to provide a comprehensive response that includes the steps and considerations for traveling to both destinations. First, I should mention the importance of planning, which includes deciding on the duration of the trip, the budget, and the type of experiences they are looking for. Next, I should provide information on how to get to each destination, including flight options and the best times to travel. Additionally, I should highlight key attractions and activities in both Thailand and California to give the customer an idea of what to expect. Finally, I should touch on practical considerations such as visa requirements for Thailand, accommodation options, and any cultural tips to enhance their travel experience.</hyde>\n",
       "\n",
       "To go on vacation in Thailand and California, start by planning your trip, including setting a budget and deciding on the duration; for Thailand, consider flying into Bangkok or Phuket and explore attractions like the Grand Palace and beaches, while in California, fly into Los Angeles or San Francisco to visit iconic sites such as Hollywood and Yosemite National Park, and don’t forget to check visa requirements for Thailand and book accommodations in advance for both destinations."
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "queries = hyde_chain.invoke({\"question\": \"How do go on vacation in thailand and in California?\"})\n",
    "display_markdown(Markdown(queries))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "841d8d70-8801-40f6-bc29-67d10de85228",
   "metadata": {},
   "source": [
    "In this section we demonstrated the possiblity of augmented the original message to produce stronger results. Naturally, this LLM-driven approach requires an additional inference, which introduces some additional latency.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c2aa6a-d144-4b0b-a5ca-7c8cd2ce2d4e",
   "metadata": {},
   "source": [
    "## Model answer generation\n",
    "\n",
    "In most RAG pipelines, the number of documents shown to the model is driven by the retrieval mechanism. This generally returns up to some static number of documents provided they meeting the necessary similarity treshold. Often, this results in irrelevant documents being sent to the model for inference. Although we can easily intruct the model to ignore irrelevant documents, it is often useful for the model to explicitly call-out the documents it did use. Furthermore, many lines of research have demonstrated the effectiveness of enabling the model to correct itself. In both cases, we make an additional call to the model once an initial answer is generated in order to improve the output for the end-user. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da954d0-e9a9-4d66-a8bd-e385d175ed6a",
   "metadata": {},
   "source": [
    "### Citation\n",
    "\n",
    "We generate an output with `answer` and `docs` keys. `docs` contains a list of Langchain `Document` objects. These are the documents the model has picked as being relevant to answering the original query. Although the documents are currently returned with title and summaries, these keys are part of a `metadata` attribute letting you determine any number of field that may be relevant to be used by your application such as author, source URL, etc... \n",
    "\n",
    "We define the system prompt to generate the model answer. Note that this is a simple template that can be further augmented with additional sections better describing our task and intended output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "42545ca3-8139-4e53-b30b-351d64687ed7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:11:26.233377Z",
     "start_time": "2025-02-13T20:11:26.231305Z"
    }
   },
   "outputs": [],
   "source": [
    "citation_system_prompt = \"\"\"You're a helpful AI assistant. Given a user question and some article snippets, answer the user question. \n",
    "If none of the articles answer the question, just say you don't know.\n",
    "\n",
    "Here are the articles: {context}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d6282c2-7963-4b49-b3d9-5897d9cf5204",
   "metadata": {},
   "source": [
    "This prompt is past as part the broader chat template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "30c8d42b-b30c-4ed2-888b-3cf92bb2095b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:11:30.988556Z",
     "start_time": "2025-02-13T20:11:30.984968Z"
    }
   },
   "outputs": [],
   "source": [
    "citation_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", citation_system_prompt),\n",
    "        (\"human\", \"Here is the customer's question: <question>{question}</question> How do you answer to the instructions?\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "answer_generator = citation_prompt | nova_pro | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b920107b-6736-4828-9be1-5fc4a72dfd3c",
   "metadata": {},
   "source": [
    "Lets use the `WikipediaRetriever` allowing us to interact with the Wikipedia API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "54740d6b-f8ee-4a3c-a3b5-0bbcf264fad3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:10.434209Z",
     "start_time": "2025-02-13T20:12:10.383291Z"
    }
   },
   "outputs": [],
   "source": [
    "wiki = WikipediaRetriever(top_k_results=6, doc_content_chars_max=2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb40110-0144-489f-ab17-32b610c8eb8f",
   "metadata": {},
   "source": [
    "The `format_docs` helper function is used to format the documents returned by the retriever to make them more friendly to the model. We supply the document's title and summary snippet. At the end, we pass the function to a child of Lanchain's `Runnable` class. This simply enables us to call the function with a standard API (invoke, batch, stream, transform and compose). Many object in Langchain implement this interface including `BaseModel`. \n",
    "\n",
    "To demonstrate the power of citations, we also append an additional obviously irrelevant document to the formatted documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "954c48d0-df8c-4f95-be9e-51dae41d0e5c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:12.188942Z",
     "start_time": "2025-02-13T20:12:12.186041Z"
    }
   },
   "outputs": [],
   "source": [
    "def format_docs(docs: List[Document]) -> str:\n",
    "    \"\"\"Convert Documents to a single string.:\"\"\"\n",
    "    formatted = [\n",
    "        f\"Article Title: {doc.metadata['title']}\\nArticle Snippet: {doc.page_content}\"\n",
    "        for doc in docs\n",
    "    ]\n",
    "    formatted.append(\"Article Title: This is an irrelevant document \\\n",
    "    Article Snippet: The document is most irrelevant.\")\n",
    "    return \"\\n\\n\" + \"\\n\\n\".join(formatted)\n",
    "\n",
    "\n",
    "format = itemgetter(\"docs\") | RunnableLambda(format_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9539ea88-f505-47e2-a6af-9153655077ce",
   "metadata": {},
   "source": [
    "We define a chain as `RunnableParallel` object, which is an extention of `Runnable` that runs a mapping of Runnables in parallel, and returns a mapping of their outputs. We set the question property using `RunnablePassthrough`. This passes the input unchanged. Then, we assign values to keys in the prompt templates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "18dab078-a3ca-4873-9534-65cb48d731cc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:15.796567Z",
     "start_time": "2025-02-13T20:12:15.790509Z"
    }
   },
   "outputs": [],
   "source": [
    "citation_chain = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format)\n",
    "    .assign(answer=answer_generator)\n",
    "    .pick([\"answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd5fb5bb-c78f-4ba9-9693-764f3e1c39a9",
   "metadata": {},
   "source": [
    "When invoking the chain, it returns the original answer and the documents used for generation. Notice that some documents are relevant to the final answer and some are not. We can address this challenge with further LLM or metadata document filtering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "ea89c0af-063e-4c40-bf6b-961b5d920c47",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:28.882199Z",
     "start_time": "2025-02-13T20:12:19.915919Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'To go on vacation in Thailand and California, you would typically follow these steps:\\n\\n### Planning Your Trip\\n\\n1. **Research**:\\n   - **Destinations**: Identify the specific places you want to visit in both Thailand and California. For Thailand, popular destinations include Bangkok, Phuket, Chiang Mai, and Koh Samui. In California, consider Los Angeles, San Francisco, San Diego, and Napa Valley.\\n   - **Activities**: Look into activities and attractions you’d like to experience. Thailand offers cultural experiences, beaches, and temples, while California provides a mix of urban adventures, beaches, wine tours, and national parks.\\n\\n2. **Budget**:\\n   - Estimate your budget for flights, accommodations, food, activities, and miscellaneous expenses.\\n\\n3. **Timing**:\\n   - Check the best times to visit both locations. Thailand’s peak season is from November to February, while California’s varies by region but generally, late spring to early fall is popular.\\n\\n### Booking\\n\\n1. **Flights**:\\n   - Book international flights to Thailand and domestic flights within the US if needed. Consider layovers and travel time.\\n\\n2. **Accommodations**:\\n   - Reserve hotels, resorts, or other lodging options in both Thailand and California. Websites like Booking.com, Airbnb, and Hotels.com can be useful.\\n\\n3. **Transportation**:\\n   - Arrange for local transportation. In Thailand, you might use tuk-tuks, Grab (similar to Uber), or public transport. In California, renting a car might be convenient for traveling between cities.\\n\\n### Documentation\\n\\n1. **Passports and Visas**:\\n   - Ensure your passport is valid for at least six months beyond your planned stay. Check visa requirements for Thailand; many nationalities can enter for a short stay without a visa, but it’s good to confirm.\\n\\n2. **Vaccinations**:\\n   - Check if any vaccinations are recommended or required for travel to Thailand.\\n\\n### Packing\\n\\n1. **Clothing**:\\n   - Pack accordingly for the climates. Thailand is tropical, so light clothing is essential. California can vary, so pack layers.\\n\\n2. **Essentials**:\\n   - Don’t forget travel documents, medications, adapters for electronic devices, and any specific items you might need.\\n\\n### On the Day of Travel\\n\\n1. **Airport**:\\n   - Arrive at the airport with plenty of time to check-in, go through security, and board your flight.\\n\\n### Enjoy Your Vacation\\n\\nOnce you arrive, follow your planned itinerary but also leave room for spontaneous adventures. Enjoy the unique experiences each location has to offer!\\n\\nIf you need more specific details or have other questions, feel free to ask!',\n",
       " 'docs': [Document(metadata={'title': 'Belinda Carlisle', 'summary': 'Belinda Jo Carlisle ( KAR-lyle; born August 17, 1958) is an American singer and songwriter. She gained fame as the lead vocalist of the Go-Go\\'s, one of the most successful all-female rock bands of all time, and went on to have a prolific career as a solo artist.\\nRaised in Southern California, Carlisle was the lead vocalist of the Go-Go\\'s, which she co-founded in 1978. With their chart-topping debut studio album Beauty and the Beat in 1981, the group helped popularize new wave music in the United States. The Go-Go\\'s have sold over seven million records worldwide.\\nAfter the break-up of the Go-Go\\'s in 1985, Carlisle went on to have a successful solo career with radio hits such as \"Mad About You\", \"I Get Weak\", \"Circle in the Sand\", \"Leave a Light On\", \"Summer Rain\", and \"Heaven Is a Place on Earth\". The Go-Go\\'s reformed in 1999; Carlisle maintained her solo career and performed with the band until its disbandment in 2022.\\nCarlisle\\'s autobiography, Lips Unsealed, published in June 2010, was a New York Times Best Seller and received favorable reviews. In 1999, Carlisle was ranked No. 76 with the Go-Go\\'s in VH1\\'s 100 Greatest Women of Rock & Roll. In 2011, Carlisle, as a member of the Go-Go\\'s, received a star on the Hollywood Walk of Fame. She and the band were inducted into the Rock and Roll Hall of Fame in 2021 and the California Hall of Fame in 2024.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Belinda_Carlisle'}, page_content='Belinda Jo Carlisle ( KAR-lyle; born August 17, 1958) is an American singer and songwriter. She gained fame as the lead vocalist of the Go-Go\\'s, one of the most successful all-female rock bands of all time, and went on to have a prolific career as a solo artist.\\nRaised in Southern California, Carlisle was the lead vocalist of the Go-Go\\'s, which she co-founded in 1978. With their chart-topping debut studio album Beauty and the Beat in 1981, the group helped popularize new wave music in the United States. The Go-Go\\'s have sold over seven million records worldwide.\\nAfter the break-up of the Go-Go\\'s in 1985, Carlisle went on to have a successful solo career with radio hits such as \"Mad About You\", \"I Get Weak\", \"Circle in the Sand\", \"Leave a Light On\", \"Summer Rain\", and \"Heaven Is a Place on Earth\". The Go-Go\\'s reformed in 1999; Carlisle maintained her solo career and performed with the band until its disbandment in 2022.\\nCarlisle\\'s autobiography, Lips Unsealed, published in June 2010, was a New York Times Best Seller and received favorable reviews. In 1999, Carlisle was ranked No. 76 with the Go-Go\\'s in VH1\\'s 100 Greatest Women of Rock & Roll. In 2011, Carlisle, as a member of the Go-Go\\'s, received a star on the Hollywood Walk of Fame. She and the band were inducted into the Rock and Roll Hall of Fame in 2021 and the California Hall of Fame in 2024.\\n\\n\\n== Early life and education ==\\n\\nBelinda Jo Carlisle was born in Hollywood, Los Angeles, California, on August 17, 1958, to Harold Carlisle, a gas station employee, and his wife, Joanne (née Thompson), a homemaker. Her mother met her father, who was 20 years her senior, at age 18, and Carlisle was born nine months later. She was named after her mother\\'s favorite film, Johnny Belinda (1948). Carlisle was the first of seven siblings; she has three brothers and three sisters. When she was five years old, Carlisle\\'s father abandoned their family, and she has stated that she spent most of her childhood impoverished. As a teena'),\n",
       "  Document(metadata={'title': 'List of minimum annual leave by country', 'summary': 'In the majority of nations, including all industrialised nations except the United States, advances in employee relations have seen the introduction of statutory agreements for minimum employee leave from work—that is the amount of entitlement to paid vacation and public holidays. Companies may offer contractually more time. Companies and the law may also differ as to whether public holidays are counted as part of the minimum leave.\\nDisparities in national minimums are still subject of debate regarding work-life balance and perceived differences between nations. These numbers usually refer to full-time employment – part-time workers may get a reduced number of days. In most countries, public holidays are paid and usually not considered part of the annual leave. Also, in most countries there are additional paid leave benefits such as parental leave and sick leave that are not listed here.', 'source': 'https://en.wikipedia.org/wiki/List_of_minimum_annual_leave_by_country'}, page_content='In the majority of nations, including all industrialised nations except the United States, advances in employee relations have seen the introduction of statutory agreements for minimum employee leave from work—that is the amount of entitlement to paid vacation and public holidays. Companies may offer contractually more time. Companies and the law may also differ as to whether public holidays are counted as part of the minimum leave.\\nDisparities in national minimums are still subject of debate regarding work-life balance and perceived differences between nations. These numbers usually refer to full-time employment – part-time workers may get a reduced number of days. In most countries, public holidays are paid and usually not considered part of the annual leave. Also, in most countries there are additional paid leave benefits such as parental leave and sick leave that are not listed here.\\n\\n\\n== Methodology ==\\nFor the purpose of comparison, the paid vacation column has been normalised to a five-day workweek. For instance, a calendar month is divided by seven and multiplied by five, while a six-day workweek day is divided by six and multiplied by five. The paid vacation column gives the minimum mandatory vacation days for an employee who has one year of service with the same employer.\\nIn some countries, the public holidays are strictly bound to the calendar dates, so if they happen on Saturday or Sunday, they are \"lost\" for that year. As a result, the average number of paid extra free days can be lower than the table shows. For example, in the Czech Republic, where the official number of paid public holidays is 13, the average number of public holidays during working days in the years 2000–2016 was only 8.9 days. In other countries, such as the United Kingdom and the United States, the public holidays which would fall on Saturday or Sunday are moved to the nearest Monday or Friday.\\n\\n\\n== Countries ==\\n\\n\\n== See also ==\\nAnnual leave\\nLong service leave\\nParental leave\\nList of'),\n",
       "  Document(metadata={'title': 'The Hangover Part II', 'summary': 'The Hangover Part II is a 2011 American comedy film produced by Legendary Pictures and distributed by Warner Bros. Pictures. The sequel to the 2009 film The Hangover and the second installment in The Hangover trilogy, the film was directed by Todd Phillips, who co-wrote the script with Craig Mazin and Scot Armstrong, and stars Bradley Cooper, Ed Helms, Zach Galifianakis, Ken Jeong, Jeffrey Tambor, Justin Bartha, and Paul Giamatti.\\nIt tells the story of Phil, Stu, Alan, and Doug, as they travel to Thailand. After the bachelor party in Las Vegas, Stu takes no chances and opts for a safe, subdued pre-wedding brunch. Things do not go as planned, resulting in another bad hangover with no memories of the previous night.\\nDevelopment began in April 2009, two months before The Hangover was released. The principal actors were cast in March 2010 to reprise their roles from the first film. Production began in October 2010, in Ontario, California, before moving on location in Thailand. The film was released on May 26, 2011, and became the eighth-highest-grossing film of 2011 and the highest-grossing R-rated comedy during its theatrical run, and received mixed reviews.\\nA third installment, The Hangover Part III, was released on May 24, 2013.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/The_Hangover_Part_II'}, page_content=\"The Hangover Part II is a 2011 American comedy film produced by Legendary Pictures and distributed by Warner Bros. Pictures. The sequel to the 2009 film The Hangover and the second installment in The Hangover trilogy, the film was directed by Todd Phillips, who co-wrote the script with Craig Mazin and Scot Armstrong, and stars Bradley Cooper, Ed Helms, Zach Galifianakis, Ken Jeong, Jeffrey Tambor, Justin Bartha, and Paul Giamatti.\\nIt tells the story of Phil, Stu, Alan, and Doug, as they travel to Thailand. After the bachelor party in Las Vegas, Stu takes no chances and opts for a safe, subdued pre-wedding brunch. Things do not go as planned, resulting in another bad hangover with no memories of the previous night.\\nDevelopment began in April 2009, two months before The Hangover was released. The principal actors were cast in March 2010 to reprise their roles from the first film. Production began in October 2010, in Ontario, California, before moving on location in Thailand. The film was released on May 26, 2011, and became the eighth-highest-grossing film of 2011 and the highest-grossing R-rated comedy during its theatrical run, and received mixed reviews.\\nA third installment, The Hangover Part III, was released on May 24, 2013.\\n\\n\\n== Plot ==\\nStu Price will travel to Thailand for his upcoming wedding to Lauren, his fiancée. To avoid what happened in Las Vegas, Stu does not allow his three best friends, Doug Billings, Phil Wenneck and Alan Garner to throw him a bachelor party. He instead hosts his bachelor party at IHOP with Phil and Doug.\\nTracy gets Doug to convince Stu to allow her brother Alan to accompany them to Thailand. Stu allows his friends (including Alan), Tracy and Phil's wife Stephanie to go with him. At the airport, they are joined by Lauren's 16-year-old brother Teddy, a Stanford University scholar.\\nAt the rehearsal dinner, Lauren's father Fong reveals his disapproval of Stu during a toast. Later that night, Stu joins Phil, Doug, Alan and Teddy for a bee\"),\n",
       "  Document(metadata={'title': 'Patrick Schwarzenegger', 'summary': 'Patrick Arnold Shriver Schwarzenegger (born September 18, 1993) is an American actor. He is the son of Arnold Schwarzenegger and Maria Shriver. He began his career playing minor roles in the early 2000s, and has since starred in the television series The Staircase (2022), American Sports Story (2024), and the third season of The White Lotus (2025).\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Patrick_Schwarzenegger'}, page_content=\"Patrick Arnold Shriver Schwarzenegger (born September 18, 1993) is an American actor. He is the son of Arnold Schwarzenegger and Maria Shriver. He began his career playing minor roles in the early 2000s, and has since starred in the television series The Staircase (2022), American Sports Story (2024), and the third season of The White Lotus (2025).\\n\\n\\n== Early life and family ==\\nSchwarzenegger was born on September 18, 1993, at Providence St John's Health Center in Santa Monica, California, and raised in Los Angeles, California. He is the eldest son of Arnold Schwarzenegger (an Austrian-born bodybuilder, actor, and former governor of California) and Maria Shriver (an author and journalist). Schwarzenegger has two older sisters, Katherine and Christina, a younger brother, Christopher, and a younger paternal half-brother, Joseph Baena. He holds dual Austrian-American citizenship, speaks English and a little German, and regularly visits Austria. He is a cousin of Patrick M. Knapp Schwarzenegger.\\nWhile in high school at Brentwood School, Schwarzenegger took private acting lessons with Nancy Banks. In 2012, he matriculated at the University of Southern California, and in May 2016 received his bachelor’s degree from the USC Marshall School of Business, with a major in business administration and a minor in cinematic arts. He was part of the Lambda Chi Alpha fraternity.\\n\\n\\n== Career ==\\n\\n\\n=== Acting ===\\n\\nAt the age of 10, Schwarzenegger had a small role in the film The Benchwarmers. Throughout his childhood, he practiced acting with his father, and through his young adult years, he studied theater at USC while taking acting classes weekly at Nancy Banks Studio. At 15 years old, he started a clothing line.\\nIn the following years, Schwarzenegger had supporting roles in 2012's Stuck in Love, 2013's Grown Ups 2 and 2015's Scouts Guide to the Zombie Apocalypse. His first leading role was opposite Bella Thorne in 2018's Midnight Sun, a romantic drama about a teenage girl with a rar\"),\n",
       "  Document(metadata={'title': 'Medical tourism', 'summary': 'Medical tourism is the practice of traveling abroad to obtain medical treatment. In the past, this usually referred to those who traveled from less-developed countries to major medical centers in highly developed countries for treatment unavailable at home. However, in recent years it may equally refer to those from developed countries who travel to developing countries for lower-priced medical treatments. With differences between the medical agencies, such as the Food and Drug Administration (FDA) or the European Medicines Agency (EMA), etc., which decide whether a drug is approved in their country or region, or not, the motivation may be also for medical services unavailable or non-licensed in the home country.\\nMedical tourism most often is for surgeries (cosmetic or otherwise) or similar treatments, though people also travel for dental tourism or fertility tourism. People with rare conditions may travel to countries where the treatment is better understood. However, almost all types of health care are available, including psychiatry, alternative medicine, convalescent care, and even burial services.\\nHealth tourism is a wider term for travel that focuses on medical treatments and the use of healthcare services. It covers a wide field of health-oriented tourism ranging from preventive and health-conductive treatment to rehabilitational and curative forms of travel. Wellness tourism is a related field.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Medical_tourism'}, page_content='Medical tourism is the practice of traveling abroad to obtain medical treatment. In the past, this usually referred to those who traveled from less-developed countries to major medical centers in highly developed countries for treatment unavailable at home. However, in recent years it may equally refer to those from developed countries who travel to developing countries for lower-priced medical treatments. With differences between the medical agencies, such as the Food and Drug Administration (FDA) or the European Medicines Agency (EMA), etc., which decide whether a drug is approved in their country or region, or not, the motivation may be also for medical services unavailable or non-licensed in the home country.\\nMedical tourism most often is for surgeries (cosmetic or otherwise) or similar treatments, though people also travel for dental tourism or fertility tourism. People with rare conditions may travel to countries where the treatment is better understood. However, almost all types of health care are available, including psychiatry, alternative medicine, convalescent care, and even burial services.\\nHealth tourism is a wider term for travel that focuses on medical treatments and the use of healthcare services. It covers a wide field of health-oriented tourism ranging from preventive and health-conductive treatment to rehabilitational and curative forms of travel. Wellness tourism is a related field.\\n\\n\\n== History ==\\nThe first recorded instance of people traveling for medical treatment dates back thousands of years to when Greek pilgrims traveled from the eastern Mediterranean to a small area in the Saronic Gulf called Epidauria. This territory was the sanctuary of the healing god Asklepios.\\nSpa towns and sanitaria were early forms of medical tourism. In 18th-century Europe patients visited spas because they were places with supposedly health-giving mineral waters, treating diseases from gout to liver disorders and bronchitis.\\n\\n\\n== Description ==\\nFactors that have '),\n",
       "  Document(metadata={'title': 'Pump (album)', 'summary': 'Pump is the tenth studio album by American rock band Aerosmith. It was released on September 12, 1989, by Geffen Records. The album peaked at No. 5 on the US charts, and was certified septuple platinum by the RIAA in 1995.\\nThe album contains the hit singles \"Love in an Elevator\", \"The Other Side\", \"What It Takes\", \"Janie\\'s Got a Gun\", which all entered the Top 40 of the Hot 100. It also has certified sales of seven million copies in the U.S. to date, and is tied with its successor Get a Grip as Aerosmith\\'s second best-selling studio album in the U.S. (Toys in the Attic leads with nine million). It produced a variety of successes and \"firsts\" for the band including their first Grammy Award (\"Janie\\'s Got a Gun\"). \"Love in an Elevator\" became the first Aerosmith song to hit number one on the Mainstream Rock Tracks chart. The album was the fourth best-selling album of the year 1990.\\nIn the UK, it was the second Aerosmith album to be certified Silver (60,000 units sold) by the British Phonographic Industry, achieving this in September 1989.\\nPump was the second of three sequentially recorded Aerosmith albums to feature producer Bruce Fairbairn and engineers Mike Fraser and Ken Lomas at Little Mountain Sound Studios.\\nTwo video documentaries on the recording, Things That Go Pump in the Night and The Making of Pump, were released on VHS and LaserDisc in 1990, with the latter also released on DVD in 1997.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Pump_(album)'}, page_content='Pump is the tenth studio album by American rock band Aerosmith. It was released on September 12, 1989, by Geffen Records. The album peaked at No. 5 on the US charts, and was certified septuple platinum by the RIAA in 1995.\\nThe album contains the hit singles \"Love in an Elevator\", \"The Other Side\", \"What It Takes\", \"Janie\\'s Got a Gun\", which all entered the Top 40 of the Hot 100. It also has certified sales of seven million copies in the U.S. to date, and is tied with its successor Get a Grip as Aerosmith\\'s second best-selling studio album in the U.S. (Toys in the Attic leads with nine million). It produced a variety of successes and \"firsts\" for the band including their first Grammy Award (\"Janie\\'s Got a Gun\"). \"Love in an Elevator\" became the first Aerosmith song to hit number one on the Mainstream Rock Tracks chart. The album was the fourth best-selling album of the year 1990.\\nIn the UK, it was the second Aerosmith album to be certified Silver (60,000 units sold) by the British Phonographic Industry, achieving this in September 1989.\\nPump was the second of three sequentially recorded Aerosmith albums to feature producer Bruce Fairbairn and engineers Mike Fraser and Ken Lomas at Little Mountain Sound Studios.\\nTwo video documentaries on the recording, Things That Go Pump in the Night and The Making of Pump, were released on VHS and LaserDisc in 1990, with the latter also released on DVD in 1997.\\n\\n\\n== Production ==\\nIn December 1988, Aerosmith got together at Rik Tinory Productions in Cohasset, Massachusetts to rehearse and compose new songs, as the band members thought the isolated nature of the studio would help their creativity. Over 19 songs were written, split between an \"A-list\" with songs considered possible hits, such as \"Love in an Elevator\" and \"What It Takes\", and the \"B-list\" having songs yet to be developed such as \"Voodoo Medicine Man\". Producer Bruce Fairbairn focused on getting as many hooks on the songs as possible.\\nSome songs proposed for the album, ')]}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "citation_chain.invoke(\"How do go on vacation in thailand and in California?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb664b45-a5cf-4c7a-865c-706b96a27385",
   "metadata": {},
   "source": [
    "### Self-validation\n",
    "\n",
    "Giving the model an opportunity to correct itself has been shown to increase performance on a number of tasks. We perform self-validation and define a set of formatting rules that align with the conversational tone we expect to have from our application. We define a system prompt with this task and set of rules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "dcb47a9d-ce6b-44db-9a1e-33fbf3c9dfff",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:38.454961Z",
     "start_time": "2025-02-13T20:12:38.451796Z"
    }
   },
   "outputs": [],
   "source": [
    "valid_system_prompt = \"\"\"You are a validator and message synthesize. \n",
    "Your task is to create one coherent answer and double check the original responses to the question {question} for common mistakes, including:\n",
    "- Answer in bullet points. It should be a complete paragraph instead.\n",
    "- Inaccuracies or things that seem impossible\n",
    "\n",
    "If there are any of the above mistakes, rewrite the response. If there are no mistakes, just reproduce the original response.\n",
    "Think about your answer first before you respond. \n",
    "If some exist, put all the issues and then put your final response in <validation></validation> tags and do not include anything after.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a703b1-5978-4d26-9b20-57bc61a4e980",
   "metadata": {},
   "source": [
    "We define the prompt template with the system prompt and original model answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "005d6b62-78bb-4086-abfe-38454b2420b9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:12:44.098053Z",
     "start_time": "2025-02-13T20:12:44.096156Z"
    }
   },
   "outputs": [],
   "source": [
    "validation_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", valid_system_prompt), \n",
    "        (\"human\", \"Here is the original message produced: <orignal_message>{original}</orignal_message> How do you answer to the instructions?\")]\n",
    ")\n",
    "validation_chain = validation_prompt | nova_pro | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "686fb39d-3246-4b02-873e-593c4fef3b74",
   "metadata": {},
   "source": [
    "We invoke model, which points out obvious issues in the original document and answers with a more consistent alternative. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "f3fa985b-d328-4f1b-9989-8d1dcaecf0c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:13:06.413602Z",
     "start_time": "2025-02-13T20:13:04.060042Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The original message contains inaccuracies and is overly simplistic. Traveling from Montreal to Thailand by car is not feasible due to geographical constraints and the need to cross multiple countries and bodies of water. The only practical way to travel directly from Montreal to Thailand is by plane.\n",
       "\n",
       "<validation>\n",
       "To travel from Montreal to Thailand, the most practical and common method is by plane. There are no direct flights from Montreal to Thailand, so you would typically need to book a flight with one or two layovers, depending on your preferred airline and routing. Major airlines often fly this route, with common layover locations including cities in Europe or the Middle East. It’s important to check for the latest travel advisories, visa requirements, and flight availability before booking. Ensure you have all necessary documentation, including a passport valid for at least six months beyond your planned stay and any required visas.\n",
       "</validation>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "validation = validation_chain.invoke({\n",
    "    \"question\" : \"how to go to thailand from Montreal?\",\n",
    "    \"original\": \"1- by plane 2-by car.\",\n",
    "})\n",
    "display_markdown(Markdown(validation))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d9bebfc-9aa1-4f33-b4a2-400c7fc5ead8",
   "metadata": {},
   "source": [
    "## Putting it all together \n",
    "\n",
    "The previous components offer important primitives to build a performant RAG solution. They act as building blocks of a broader solution. We provide an example showcasing how they can be brought together in a single chain to improve response accuracy. To minimize latency and improve accuracy, we use Amazon Nova Lite for simpler tasks and Nova Pro where we need more performance.\n",
    "\n",
    "First, we create helper functions to parse the return messages for the relevant section that can be found in between tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "34226c16-cea3-4989-8ffe-40390228e68a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:13:09.613170Z",
     "start_time": "2025-02-13T20:13:09.610886Z"
    }
   },
   "outputs": [],
   "source": [
    "def parse_intent(ai_message: AIMessage) -> str:\n",
    "    \"\"\"Parse the AI message.\"\"\"\n",
    "    intent_pattern = r\"<intention>(.*?)</intention>\"\n",
    "    intent_match = re.findall(intent_pattern, ai_message.content, flags=0)\n",
    "    if intent_match:\n",
    "        return intent_match[0]\n",
    "    else:\n",
    "        return \"No intention found.\"\n",
    "\n",
    "def parse_norm_message(ai_message: AIMessage) -> str:\n",
    "    \"\"\"Parse the AI message.\"\"\"\n",
    "    norm_pattern = r\"<condensed_message>(.*?)</condensed_message>\"\n",
    "    norm_match = re.findall(norm_pattern, ai_message['question'].content, flags=0)\n",
    "    if norm_match:\n",
    "        return norm_match[0]\n",
    "    else:\n",
    "        return \"Message could not be successfully normalized.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b75e6b95-a615-408d-8824-33f79a54d83c",
   "metadata": {},
   "source": [
    "We define an end-to-end RAG chain primairly using LangChain Expression Language (LCEL), which allows us to define `Runnable` objects in success to one another. The resulting chain reuses many of the components we previously defined including intent detection with **dynamic few-shots, message normalization and citation**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "2b6b651f-3b9b-4fe3-aff2-26e5b44920e4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-13T20:15:15.295813Z",
     "start_time": "2025-02-13T20:15:06.736669Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "Certainly! Here are some steps to help you plan a vacation:\n",
       "\n",
       "### 1. **Determine Your Budget**\n",
       "   - **Set a Budget**: Decide how much you are willing to spend on your vacation. This should include travel, accommodation, food, activities, and miscellaneous expenses.\n",
       "   - **Research Costs**: Look up the average costs for flights, accommodations, and activities at your desired destination.\n",
       "\n",
       "### 2. **Choose a Destination**\n",
       "   - **Interests and Preferences**: Consider what you want to do on your vacation (relax on a beach, explore cities, go hiking, etc.).\n",
       "   - **Season and Weather**: Check the weather conditions and seasonal events at your potential destinations.\n",
       "   - **Visa and Travel Requirements**: Ensure you have the necessary documents and vaccinations required for your chosen destination.\n",
       "\n",
       "### 3. **Plan Your Itinerary**\n",
       "   - **Duration**: Decide how long you want to stay.\n",
       "   - **Activities**: Make a list of must-see attractions and activities. Balance between sightseeing and relaxation.\n",
       "   - **Transportation**: Plan how you will get around (public transport, rental car, etc.).\n",
       "\n",
       "### 4. **Book Flights and Accommodation**\n",
       "   - **Flights**: Compare prices on different airlines and book in advance for the best deals.\n",
       "   - **Accommodation**: Choose between hotels, vacation rentals, hostels, or other options based on your budget and preferences. Book early for popular destinations.\n",
       "\n",
       "### 5. **Pack Accordingly**\n",
       "   - **Clothing**: Pack based on the weather and activities planned.\n",
       "   - **Essentials**: Don’t forget travel documents, medications, chargers, and any specific items you’ll need.\n",
       "\n",
       "### 6. **Consider Travel Insurance**\n",
       "   - **Coverage**: Look into plans that cover medical emergencies, trip cancellations, and lost luggage.\n",
       "\n",
       "### 7. **Final Preparations**\n",
       "   - **Notify Banks**: Inform your bank of your travel plans to avoid issues with your credit/debit cards.\n",
       "   - **Download Apps**: Useful travel apps for navigation, translations, and booking can be helpful.\n",
       "   - **Backup Plans**: Have a backup plan for popular attractions that may be closed or overcrowded.\n",
       "\n",
       "By following these steps, you can ensure a well-organized and enjoyable vacation. Safe travels!"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rag_chain = RunnableParallel(\n",
    "    question=RunnablePassthrough(),\n",
    "    intent=few_shot_intent_detection_prompt | nova | parse_intent\n",
    ") | RunnableBranch(\n",
    "    (lambda payload: \"vacation\" == payload[\"intent\"].lower(), lambda x: (\n",
    "        RunnablePassthrough().pick([\"question\"])\n",
    "        .assign(question=norm_chain)\n",
    "        .assign(question=parse_norm_message)\n",
    "        .assign(context=lambda inputs: wiki.invoke(inputs[\"question\"]))\n",
    "        .assign(answer=answer_generator)\n",
    "        .pick([\"answer\", \"context\"])\n",
    "    )),\n",
    "    (lambda payload: \"irrelevant\" == payload[\"intent\"].lower(), lambda x: AIMessage(content=\"I am only able to answer questions about travel and vacations.\")),\n",
    "    (lambda payload: \"contact\" == payload[\"intent\"].lower(), lambda x: AIMessage(content=\"I am transfering you to an agent now...\")),\n",
    "    lambda payload: AIMessage(content=\"I am only able to answer questions about travel and vacations.\" )\n",
    ")\n",
    "\n",
    "display_markdown(Markdown(rag_chain.invoke(\"I want to know more about how to plan a vacation?\")['answer']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d4040b0-eaae-4f12-a98a-cdc777392b52",
   "metadata": {},
   "source": [
    "It is evident that latency is increased in corralation with the number calls being made in succession. Hence, it is optimal to make calls in parallel where possible to reduce overall time to execute the entire pipeline. Notice in in our example that the intent detection could be made in parallel to message normalization and citation (model inference).\n",
    "\n",
    "Additionally, it may be benifitial to modify the pipeline to include a query augmentation step for reasons described earlier in the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d52a0ce-4d5e-4f0f-bd01-326a249fa4a1",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Where RAG enables single-turn conversations where users and agents alternate sending eachother messages, agents supply the ability to the application developer to build increased complexity into the conversation flow. These applications are characterized by increase **autonomy, reactivity, proactiveness, adaptability and situatedness**. They typically have some form of validation, the ability to loop back and call external functions to improve outputs. You can dive deeper into agents in the next lab of this workshop."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77518672-412b-4ddc-a548-ebae5c3ff12e",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "\n",
    "There is no necessary clean up for this notebook."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
